import json
import csv
import copy
import os


class InputExample(object):
    """A single training/test example for token classification."""

    def __init__(self, guid, text_a, labels):
        """Constructs a InputExample.
        Args:
            guid: Unique id for the example.
            text_a: list. The words of the sequence.
            labels: (Optional) list. The labels for each word of the sequence. This should be
            specified for train and dev examples, but not for test examples.
        """
        self.guid = guid
        self.text_a = text_a
        self.labels = labels

    def __repr__(self):
        return str(self.to_json_string())

    def to_dict(self):
        """Serializes this instance to a Python dictionary."""
        output = copy.deepcopy(self.__dict__)
        return output

    def to_json_string(self):
        """Serializes this instance to a JSON string."""
        return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"


class DataProcessor(object):
    """Base class for data converters for sequence classification data sets."""

    def get_train_examples(self, data_dir):
        """Gets a collection of `InputExample`s for the train set."""
        raise NotImplementedError()

    def get_dev_examples(self, data_dir):
        """Gets a collection of `InputExample`s for the dev set."""
        raise NotImplementedError()

    def get_labels(self):
        """Gets the list of labels for this data set."""
        raise NotImplementedError()

    @classmethod
    def _read_tsv(cls, input_file, quotechar=None):
        """Reads a tab separated value file."""
        with open(input_file, "r", encoding="utf-8-sig") as f:
            reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
            lines = []
            for line in reader:
                lines.append(line)
            return lines

    @classmethod
    def _read_text(self, input_file):
        lines = []
        with open(input_file, 'r') as f:
            words = []
            labels = []
            for line in f:
                if line.startswith("-DOCSTART-") or line == "" or line == "\n":
                    if words:
                        lines.append({"words": words, "labels": labels})
                        words = []
                        labels = []
                else:
                    splits = line.split(" ")
                    words.append(splits[0])
                    if len(splits) > 1:
                        labels.append(splits[-1].replace("\n", ""))
                    else:
                        # Examples could have no label for mode = "test"
                        labels.append("O")
            if words:
                lines.append({"words": words, "labels": labels})
        return lines


class CnerProcessor(DataProcessor):
    """Processor for the chinese ner data set."""

    def get_train_examples(self, data_dir):
        """See base class."""
        return self._create_examples(self._read_text(os.path.join(data_dir, "train.char.bieo")), "train")

    def get_dev_examples(self, data_dir):
        """See base class."""
        return self._create_examples(self._read_text(os.path.join(data_dir, "dev.char.bieo")), "dev")

    def get_test_examples(self, data_dir):
        """See base class."""
        return self._create_examples(self._read_text(os.path.join(data_dir, "test.char.bieo")), "test")

    def get_labels(self):
        """See base class."""
        return ['B-assist', 'B-cellno', 'B-city', 'B-community', 'B-devzone', 'B-distance', 'B-district', 'B-floorno',
                'B-houseno', 'B-intersection', 'B-poi', 'B-prov', 'B-road', 'B-roadno', 'B-subpoi', 'B-town',
                'B-village_group',
                'E-assist', 'E-cellno', 'E-city', 'E-community', 'E-devzone', 'E-distance', 'E-district', 'E-floorno',
                'E-houseno', 'E-intersection', 'E-poi', 'E-prov', 'E-road', 'E-roadno', 'E-subpoi', 'E-town',
                'E-village_group',
                'I-assist', 'I-cellno', 'I-city', 'I-community', 'I-devzone', 'I-distance', 'I-district', 'I-floorno',
                'I-houseno', 'I-intersection', 'I-poi', 'I-prov', 'I-road', 'I-roadno', 'I-subpoi', 'I-town',
                'I-village_group',
                'O',
                'S-assist', 'S-intersection', 'S-poi', 'S-community', 'S-district'
                ]

    def _create_examples(self, lines, set_type):
        """Creates examples for the training and dev sets."""
        examples = []
        for (i, line) in enumerate(lines):
            if i == 0:
                continue
            guid = "%s-%s" % (set_type, i)
            text_a = line['words']
            # BIEO
            labels = []
            for x in line['labels']:
                labels.append(x)
            examples.append(InputExample(guid=guid, text_a=text_a, labels=labels))
        return examples


if __name__ == '__main__':
    a = CnerProcessor()
    data_dir = '/Users/lz/PycharmProjects/tianchi/diming_ner/ner_data'
    b = a.get_train_examples(data_dir)
    Labels = set()
    m = 0
    for l in b:
        Labels.update(l.to_dict()['labels'])
        n = len(l.to_dict()['text_a'])
        if n > m:
            m = n
